package dan.backend
import chisel3._
import chisel3.util._
import dan.common._
import dan.common.Consts._ 
import dan.frontend.BrUpdateInfo
import dan.frontend.BrResolInfo
import dan.frontend.FTQIdx2PC
import scala.collection.mutable.ArrayBuffer

class ExecUnitResp(val len:Int = 32) extends CoreBundle{
    val uop = new UOp()
    val data = Bits(len.W)
    val csrCmd = UInt(CSR_BITS.W)
    val csrAddr = UInt(CSR_ADDR_BITS.W)
    val tlbOp = UInt(TLB_IDX_BITS.W)
    val csrR1 = UInt(xBits.W)
    val csrR2 = UInt(xBits.W)
}

abstract class ExecUnit(
    val writesIrf: Boolean        = false,
    val bypassable: Boolean       = false,
    val alwaysBypassable: Boolean = false,
    val hasMem: Boolean           = false,
    val hasCSR: Boolean           = false,
    val hasJmpUnit: Boolean       = false,
    val hasAlu: Boolean           = false,
    val hasMul: Boolean           = false,
    val hasDiv: Boolean           = false,
    val hasCnt: Boolean           = false,
    val numStages: Int
) extends CoreModule {
    val io = IO(new Bundle{
        val fu_types  = Output(Bits(FU_BITS.W))

        val req       = Flipped(new DecoupledIO(new FuncUnitReq))

        val iresp     = new DecoupledIO(new ExecUnitResp)

        val bypass    = Output(Vec(numStages, Valid(new ExecUnitResp)))
        val brupdate  = Input(new BrUpdateInfo)

        val brinfo    = if (hasAlu) Output(new BrResolInfo) else null
        val getFtqPc  = if (hasJmpUnit) Flipped(new FTQIdx2PC) else null

        val lsu_io    = if (hasMem) Flipped(new LSUExeIO) else null
    })
    io <> DontCare
    require (bypassable || !alwaysBypassable, "[execute] an execution unit must be bypassable if it is always bypassable")

    def supportedFuncUnits = {
        OpFuncType(
            alu    = hasAlu,
            jmp    = hasJmpUnit,
            mem    = hasMem,
            muldiv = hasMul || hasDiv,  // should we split these up?
            csr    = hasCSR,
            cnt    = hasCnt
        )
    }

    def numReadPorts = {
        if (hasMem) 1 else 2
    }
}

class ALUExeUnit (
    hasJmpUnit: Boolean = false,
    hasCSR: Boolean     = false,
    hasAlu: Boolean     = true ,    // ALU is always supported
    hasMul: Boolean     = false,
    hasDiv: Boolean     = false,
    hasCnt: Boolean     = false,
    hasMem: Boolean     = false,
) extends ExecUnit (
    writesIrf        = hasAlu || hasMul || hasDiv || hasCnt,
    bypassable       = hasAlu || hasCnt,
    alwaysBypassable = (hasAlu || hasCnt) && !(hasMem || hasJmpUnit || hasMul || hasDiv /* || hasCSR */),
    hasMem           = hasMem,
    hasCSR           = hasCSR,
    hasJmpUnit       = hasJmpUnit,
    hasAlu           = hasAlu,
    hasMul           = hasMul,
    hasDiv           = hasDiv,
    hasCnt           = hasCnt,
    numStages        = if (hasAlu && hasMul) 2 else if (hasAlu) 1 else 0
) {
    val div_busy  = WireInit(false.B)

    val iresp_fu_units = ArrayBuffer[FuncUnit]()

    io.fu_types := Mux(hasAlu.B,              FU_ALU, 0.U) |
                   Mux(hasMul.B,              FU_MUL, 0.U) |
                   Mux(!div_busy && hasDiv.B, FU_DIV, 0.U) |
                   Mux(hasCnt.B,              FU_CNT, 0.U) |
                   Mux(hasCSR.B,              FU_CSR, 0.U) |
                   Mux(hasJmpUnit.B,          FU_JMP, 0.U) |
                   Mux(hasMem.B,              FU_MEM, 0.U)

    // ALU Unit -------------------------------
    var alu: ALUUnit = null
    if (hasAlu) {
        alu = Module(new ALUUnit(
            isJmpUnit = hasJmpUnit,
            isCSRUnit = hasCSR,
            stageNum = numStages
        ))
        alu.io.req.valid := (io.req.valid && (
            io.req.bits.uop.fuTypeEqual(FU_ALU) ||
            io.req.bits.uop.fuTypeEqual(FU_JMP) ||
            io.req.bits.uop.fuTypeEqual(FU_CSR)
        ))
        alu.io.req.bits.uop      := io.req.bits.uop
        alu.io.req.bits.kill     := io.req.bits.kill
        alu.io.req.bits.oprand1  := io.req.bits.oprand1
        alu.io.req.bits.oprand2  := io.req.bits.oprand2
        alu.io.brUpdate          := io.brupdate
        alu.io.resp.ready        := true.B
        io.bypass                := alu.io.bypass
        io.brinfo                := alu.io.brInfo
        if (hasJmpUnit) { alu.io.getFtqPC <> io.getFtqPc }

        iresp_fu_units += alu
    }

    // Counter Unit ---------------------------
    var cnt: CntUnit = null
    if (hasCnt) {
        cnt = Module(new CntUnit(stageNum = numStages))
        cnt.io.req <> DontCare
        cnt.io.req.valid := io.req.valid && io.req.bits.uop.fuTypeEqual(FU_CNT)
        cnt.io.req.bits.uop     := io.req.bits.uop
        cnt.io.req.bits.kill    := io.req.bits.kill
        cnt.io.req.bits.oprand1 := io.req.bits.oprand1
        cnt.io.req.bits.oprand2 := io.req.bits.oprand2
        cnt.io.brUpdate         := io.brupdate
        cnt.io.resp.ready       := true.B

        for (i <- 0 until numStages) {
            when (cnt.io.bypass(i).valid) {
                io.bypass(i) := cnt.io.bypass(i)
            }
        }

        iresp_fu_units += cnt
    }

    // Pipelined, IMul Unit ------------------
    var imul: PipelinedMulUnit = null
    if (hasMul) {
        imul = Module(new PipelinedMulUnit)
        imul.io.req <> DontCare
        imul.io.req.valid        := io.req.valid && io.req.bits.uop.fuTypeEqual(FU_MUL)
        imul.io.req.bits.uop     := io.req.bits.uop
        imul.io.req.bits.oprand1 := io.req.bits.oprand1
        imul.io.req.bits.oprand2 := io.req.bits.oprand2
        imul.io.req.bits.kill    := io.req.bits.kill
        imul.io.brUpdate         := io.brupdate
        imul.io.resp.ready       := true.B

        iresp_fu_units += imul
    }

    // Div/Rem Unit -----------------------
    var div: DivUnit = null
    val div_resp_val = WireInit(false.B)
    if (hasDiv) {
        div = Module(new DivUnit)
        div.io.req <> DontCare
        div.io.req.valid        := io.req.valid && io.req.bits.uop.fuTypeEqual(FU_DIV)
        div.io.req.bits.uop     := io.req.bits.uop
        div.io.req.bits.oprand1 := io.req.bits.oprand1
        div.io.req.bits.oprand2 := io.req.bits.oprand2
        div.io.brUpdate         := io.brupdate
        div.io.req.bits.kill    := io.req.bits.kill
        div.io.resp.ready       := !iresp_fu_units.map(_.io.resp.valid).reduce(_|_)

        div_resp_val := div.io.resp.valid
        div_busy     := !div.io.req.ready || (io.req.valid && io.req.bits.uop.fuTypeEqual(FU_DIV))

        iresp_fu_units += div
    }

    // Mem Unit --------------------------
    if (hasMem) {
        require(!hasAlu)
        require(numStages == 0)
        val maddrcalc = Module(new AddrGenUnit(stageNum = numStages))
        maddrcalc.io.req <> DontCare
        maddrcalc.io.req.valid  := io.req.valid && io.req.bits.uop.fuTypeEqual(FU_MEM)
        maddrcalc.io.req.bits   := io.req.bits
        maddrcalc.io.brUpdate   := io.brupdate
        maddrcalc.io.resp.ready := true.B

        io.lsu_io.req := maddrcalc.io.resp
        io.iresp <> io.lsu_io.iresp
    }

    // Outputs (Write Port #0)  ---------------
    if (writesIrf) {
        io.iresp.valid     := iresp_fu_units.map(_.io.resp.valid).reduce(_|_)
        io.iresp.bits.uop  := PriorityMux(iresp_fu_units.map(f =>
            (f.io.resp.valid, f.io.resp.bits.uop)).toSeq
        )
        io.iresp.bits.data := PriorityMux(iresp_fu_units.map(f =>
            (f.io.resp.valid, f.io.resp.bits.data)).toSeq
        )

        if (hasCSR) {
            io.iresp.bits.csrAddr := RegNext(alu.io.imm)
            io.iresp.bits.csrCmd  := alu.io.resp.bits.uop.ctrl.csrCmd
            io.iresp.bits.tlbOp   := alu.io.resp.bits.uop.tlbOp
            io.iresp.bits.csrR1   := alu.io.resp.bits.oprand1
            io.iresp.bits.csrR2   := alu.io.resp.bits.oprand2
        }
    }
    assert ((PopCount(iresp_fu_units.map(_.io.resp.valid)) <= 1.U && !div_resp_val) ||
            (PopCount(iresp_fu_units.map(_.io.resp.valid)) <= 2.U &&  div_resp_val),
        "Multiple functional units are fighting over the write port.")
}

/////////////////////////////////////

class ExecUnits extends HasCoreParam {
    private val exe_units = ArrayBuffer[ExecUnit]()

    def length = exe_units.length

    def apply(n: Int) = exe_units(n)

    def map[T](f: ExecUnit => T) = {
        exe_units.map(f)
    }

    def withFilter(f: ExecUnit => Boolean) = {
        exe_units.withFilter(f)
    }

    def foreach[U](f: ExecUnit => U) = {
        exe_units.foreach(f)
    }

    def zipWithIndex = {
        exe_units.zipWithIndex
    }

    def indexWhere(f: ExecUnit => Boolean) = {
        exe_units.indexWhere(f)
    }

    def count(f: ExecUnit => Boolean) = {
        exe_units.count(f)
    }

    lazy val memory_units = {
        exe_units.filter(_.hasMem)
    }

    lazy val alu_units = {
        exe_units.filter(_.hasAlu)
    }

    lazy val csr_unit = {
        require(exe_units.count(_.hasCSR) == 1)
        exe_units.find(_.hasCSR).get
    }

    lazy val jmp_unit_idx = {
        exe_units.indexWhere(_.hasJmpUnit)
    }

    val int_width = issueParams.find(_.iqType == IQT_INT.litValue).get.issueWidth
    val mem_width = issueParams.find(_.iqType == IQT_MEM.litValue).get.issueWidth

    for (w <- 0 until mem_width) {
        val memExeUnit = Module(new ALUExeUnit(
            hasAlu = false,
            hasMem = true
        ))
        exe_units += memExeUnit
    }
    for (w <- 0 until int_width) {
        def is_nth(n: Int): Boolean = w == (n % int_width)

        val alu_exe_unit = Module(new ALUExeUnit(
            hasJmpUnit = is_nth(0),
            hasCSR = is_nth(1),
            hasCnt = is_nth(1),
            hasMul = is_nth(2),
            hasDiv = is_nth(2),
        ))
        exe_units += alu_exe_unit
    }

    val numWritePorts       = exe_units.count(_.writesIrf)
    val numTotalBypassPorts = exe_units.withFilter(_.bypassable).map(_.numStages).sum

    val bypassable_write_port_mask = exe_units.withFilter(_.writesIrf).map(_.bypassable)
}

